#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <math.h>
#include <iostream>
#include <mpi.h>
#include <math.h>
#include <time.h>

#define MAX_LINE_SIZE 128   //fgets函数读取的最大字节数,作为缓冲

int procID, procNum, line;
double startTime, endTime, totalTime;

void sparse_add(int **J0, double **S0, int row, int col, double val);
void sparse_compress(int **J0, double **S0, int *I, int *J, double *S, int dim);
double sparse_entry(int i, int j, int *I, int *J, double *S, int dim);
//void SOR(int *I, int *J, double *S, double *rhs, double *x, int dim, double w);

int main(int argc, char* argv[])
{
    int i,j,m,n,l;
    int *I,*J,**J0,flag;
    double val,w;
    double *S,**S0;
    double *rhs, *solution;
    int procID, procNum, line;

    /// this is for the smallest sparse matrix
    char* spmat_data_file = argv[1]; // "matrix/spmat7585.csv";
    char* rhs_data_file = argv[2]; // "matrix/rhs7585.csv";
    int dim = 119425;        // 稀疏矩阵信息，这里偷懒假设事先知道！
    int nnz = 825737;

    /// 
    int maxcouple = 20;   // 每行最多有几个非零元？ 给一个适当大的数，请保证数组不越界的同时内存够用!
    char dummy, strLine[MAX_LINE_SIZE];  // buffer for reading the data files
    
    /*
     * step 1. read in J0 and S0 from data file
     */
    J0 = (int**)calloc(dim, sizeof(int*));
    S0 = (double**)calloc(dim, sizeof(double*));
    for(i=0 ; i < dim; i++){
	      J0[i] = (int*) calloc(maxcouple,sizeof(int));
	      S0[i] = (double*) calloc(maxcouple,sizeof(double));
	      for(j = 0; j < maxcouple; j++){
	          J0[i][j] = 0;
	          S0[i][j] = 0.0;
	      }
    }

    // 打开稀疏矩阵文件 -- 假设格式知道
    FILE* fin = fopen("data119425A.txt", "r");
    if (fin != NULL)
	      printf("稀疏矩阵数据文件打开成功！！\n");
    else{
	      printf("稀疏矩阵数据文件打开失败，请检查！！\n");  
        return 0;
    }
    for (l = 0; l < nnz; l++){      // 开始读取数据
	      fscanf(fin, "%d%c%d%c%lf", &i, &dummy, &j, &dummy, &val);
	      sparse_add(J0, S0, i-1, j-1, val);   // please remind the -1 for index
    }
    fclose(fin);
    
    /*
     * step 2. read in the rhs from dat file
    */
    rhs = (double*) calloc(dim, sizeof(double));
    fin = fopen("data119425b.txt", "r");
    if(fin != NULL)
	      printf("右端项数据文件打开成功！！\n");
    else{
	      printf("右端项数据文件打开失败，请检查！！\n");	return 0;
    }
    for(i = 0; i < dim; i++){
	      fscanf(fin, "%le", &val);
	      rhs[i] = val;
    }
    fclose(fin);

    /*
     * step 3: (I,J,S) is the real sparse matrix data: row compressed format
     */
    I = (int*) calloc(dim+1,sizeof(int));
    J = (int*) calloc(nnz, sizeof(int));
    S = (double*)calloc(nnz, sizeof(double));
    sparse_compress(J0,S0,I,J,S,dim);

    // Finalizing the temporal data and free the memory now
    for(i = 0; i < dim; i++){
	      free(J0[i]);
	      free(S0[i]);
    }
    free(J0); 
    free(S0);
    printf("Program Finished! There are %d nonzero entries, %d.\n",nnz,I[dim]);
    
    x = (double*) calloc(dim, sizeof(double));
    
    // SOR
    // set w ////////////////////////////////////////////////////   <-----------
    w = 1;
    //SOR(I, J, S, rhs, solution, dim, w);
    MPI_Status status;
    MPI_Init(&argc, &argv);
    MPI_Comm_rank(MPI_COMM_WORLD, &procID);
    MPI_Comm_size(MPI_COMM_WORLD, 3);
    double err = 1.0,res;

while(err > 1e-12){
    if(procID == 0)
    {      
        for(i = 1; i < 3; i++)
        {
            MPI_Send(A, dim * dim, MPI_INT, i, 0, MPI_COMM_WORLD);
        }
        for(i = 1; i < 3; i++)
        {
            MPI_Send(b, dim, MPI_INT, i, 1, MPI_COMM_WORLD);
        }
        for(k = 1; k < 3; k++)
        {
            MPI_Recv(ans, line * Length, MPI_INT, k, 3, MPI_COMM_WORLD, &status);
            // for(i = 0; i < line; i++)
            // {
            //     for( j = 0; j < Length; j++)
            //     {
            //         C[((k - 1) * line + i) * Length + j] = ans[i * Length + j];
            //     }
            // }
            for(i = 0; i < dim; i++){
                res = rhs[i];
                for(l = I[i] + 1; l < I[i+1]; l++){ 
                    res -= x[J[l]]*S[l];
                }
                x[i] =(1-w)*x[i] + w*res/S[I[i]];
            }
            //calculate the norm of the residual
            err = 0.0;
            for(i = 0; i < dim; i++){
                res = rhs[i];
                for(l = I[i]; l < I[i+1]; l++)
	                  res -= x[J[l]]*S[l];
                err += res*res;
            }
            err = sqrt(err);
        }
        for(i = (procNum - 1) * line; i < Length; i++)
        {
            for(j = 0; j < Length; j++)
            {
                temp = 0;
                for(k = 0; k < Length; k++)
                {
                    temp += A[i * Length + k] * B[k * Length + j];
                }
                C[i * Length + j] = temp;
            }
        }
    }
    else
    {
        MPI_Recv(B, Length * Length, MPI_INT, 0, 0, MPI_COMM_WORLD, &status);
        MPI_Recv(buffer, Length * line, MPI_INT, 0, 1, MPI_COMM_WORLD, &status);
        for(i = 0; i < line; i++)
        {
            for(j = 0; j < Length; j++)
            {
                temp = 0;
                for(k = 0; k < Length; k++)
                {
                    temp += buffer[i * Length + k] * B[k * Length + j];
                }
                ans[i * Length + j] = temp;
            }
        }
        MPI_Send(ans, line * Length, MPI_INT, 0, 3, MPI_COMM_WORLD);
    }
}
    while(err > 1e-12){
        for(i = 0; i < dim; i++){
            res = rhs[i];
            for(l = I[i] + 1; l < I[i+1]; l++){ 
                res -= x[J[l]]*S[l];
            }
            x[i] =(1-w)*x[i] + w*res/S[I[i]];
        }
        //calculate the norm of the residual
        err = 0.0;
        for(i = 0; i < dim; i++){
            res = rhs[i];
            for(l = I[i]; l < I[i+1]; l++)
	              res -= x[J[l]]*S[l];
            err += res*res;
        }
        err = sqrt(err);
    }

    
    for(i = 0; i < dim; ++i)
    {
        std::cout << solution[i] << std::endl;
    }       

    /// Final step: free the memory for linear system:
    free(solution);  free(rhs);
    free(I);    free(J);    free(S);

    MPI_Finalize();
    //printf("Program Finished!\n");
    return 1;
}

void sparse_add(int **J0, double **S0, int row, int col, double val)
{
    int pos;
    // find if value col is in the row's row of matrix J
    for(pos=1; pos <= J0[row][0]; pos++){
        if(J0[row][pos] == col) break;
    }
    // whether J exist or not, i is the current position to insert
    // since S is initialized to be zero, add on is safe, also the rewrite of
    // J[row][i] is also safe!! This implementation is wonderful!
    J0[row][pos] = col;
    S0[row][pos] += val;
    J0[row][0] = (pos > J0[row][0])?pos:J0[row][0];
}

/**
 * treat the sparsity pattern from matrix J0 and S0 the row compress
 * storage format is gotten : I,J,S we assume the memory is reserved by
 * calling functions.  for ease of matrix iteration method, the diagnal
 * entries are put at the beginning of each row.
 */
void sparse_compress(int **J0, double **S0, int *I, int *J, double *S, int dim)
{
    int i,j,pos;
    // fill vector I first, whose first element must be 0
    I[0] = 0;
    for(i = 1; i <= dim; i++){
        I[i] = I[i-1] + J0[i-1][0];
    }
    // then run over J0 and S0 to fill J and S
    for(i = 0; i < dim; i++){
        for(j = 1,pos = 1; j <= J0[i][0]; j++){
            // remember that pos=0 is for the diagnal element
            if(J0[i][j] == i){
	              J[I[i]] = i;
	              S[I[i]] = S0[i][j];
            }
            else{
              	J[I[i]+pos] = J0[i][j];
	              S[I[i]+pos] = S0[i][j];
	              pos++;
            } 
        } 
    }
}

/**
 * get the specific entry value on position (i,j) 
 */
double sparse_entry(int i, int j, int *I, int *J, double *S, int dim)
{
    // for ease, we assume that i and j are not beyond dim
    // but for purpose of later use, we place dim here
    int pos;
    if (i==j){ // the diagnal element is easy to get!
        pos = I[i];
        return S[pos];
    }
    // otherwise the other element
    for(pos = I[i]+1; pos < I[i+1]; pos++)
        if(J[pos] == j) break;
    if(pos == I[i+1])
        return 0.0; // (i,j) are not int the pattern, so just return 0.0
    else
        return S[pos];  
}

/**
 * Solving linear system using various sparse matrix solver.
 * Currently, SOR works well for elliptic problem.
 */
// void SOR(int *I, int *J, double *S, double *rhs, double *x, int dim, double w){
//     double err = 1.0,res;
//     int iter = 1,i,j,l;
//     // double *x_old = (double *) calloc(dim,sizeof(double));
//     while(err > 1e-12){
//         // this is SOR iteration:
//         iter++;
//         for(i = 0; i < dim; i++){
//             res = rhs[i];
//             for(l = I[i] + 1; l < I[i+1]; l++){ 
//                 res -= x[J[l]]*S[l];
//             }
//             x[i] =(1-w)*x[i] + w*res/S[I[i]];
//         }
//         //calculate the norm of the residual
//         err = 0.0;
//         for(i = 0; i < dim; i++){
//             res = rhs[i];
//             for(l = I[i]; l < I[i+1]; l++)
// 	              res -= x[J[l]]*S[l];
//             err += res*res;
//         }
//         err = sqrt(err);
//         //  if((iter++ % 100) == 0)
//         //	printf("=== Iteration %d: the error is %le ===\n",iter-1, err);
//     }
//     printf("The SOR iterative solver finished at step %d with residual = %le.\n",iter-1, err);
// }
